import cv2
from net.yolo_face_detect import YOLO_face
class FaceDetect():
    def __init__(self):
        self.yolo_face = YOLO_face('weights/yolo_face.onnx')
        pass
    def detect(self,frame):
        det_bboxes,def_conf,def_classid,landmarks = self.yolo_face.detect(frame)
        for box in det_bboxes:
            # print(box)
            x,y,w,h = box.astype('int')
            cv2.rectangle(frame,(x,y),(x+w,y+h),color=(0,255,0),thickness=1)
            face_img = frame[x:x+w,y:y+h]
        return face_img

if __name__ == '__main__':
    frame = cv2.imread('images/test/db/12_Group_Group_12_Group_Group_12_728.jpg')
    face_detect = FaceDetect()
    face_img = face_detect.detect(frame)
    cv2.imshow("detect",frame)
    cv2.imshow('face',face_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()